package ws

import (
	"context"
	"encoding/json"
	"io"
	"strings"
	"time"
	"unicode/utf8"

	"device-admin/config/mysql"
	"device-admin/internal/app/grpc/protoc/devicec"
	"device-admin/internal/app/grpc/service/terminal"
	"device-admin/internal/dal"
	"device-admin/internal/dao"
	"device-admin/internal/models"

	"192.168.1.75/go-pkg/helper"
	"192.168.1.75/go-pkg/logx"
	"github.com/gorilla/websocket"
	"google.golang.org/grpc/metadata"
	"google.golang.org/protobuf/proto"
)

// errPack 错误信息
func errPack(msg string) []byte {
	var ret struct {
		Code int64  `json:"code"`
		Msg  string `json:"msg"`
	}
	ret.Code = -1
	ret.Msg = msg
	jsonBytes, _ := json.Marshal(ret)
	return jsonBytes
}

// WebShell95Handle 堡垒机websocket客户端消息处理
func WebShell95Handle(conn *websocket.Conn, robotID, dstDir, fileName string, fileSize int64, fileMD5, userID string) {
	// 设备id不能为空
	if robotID == "" {
		conn.WriteMessage(1, errPack("设备id不能为空"))
	}
	// 设备id有效性判断
	var count int64
	mysql.NewDB().Model(&models.Device{}).Where("robot_id = ?", robotID).Count(&count)
	if count <= 0 {
		conn.WriteMessage(1, errPack("设备不存在"))
		return
	}
	// 如果上传目录不为
	{
		if dstDir != "" {
			// 文件名称
			if fileName == "" {
				conn.WriteMessage(1, errPack("文件名称不能为空"))
				return
			}
			if fileSize <= 0 {
				conn.WriteMessage(1, errPack("文件大小不能为空"))
				return
			}
		}
	}
	sessionID := helper.RandString(32)
	// 获取终端所在的服务ip,建立grpc连接，以进行数据中转
	cli, grpcConn, err := terminal.SessionServerGrpcClient(robotID, "webshell95")
	if grpcConn != nil {
		defer grpcConn.Close()
	}
	if err != nil {
		conn.WriteMessage(1, []byte(err.Error()))
		return
	}
	// webshell cli
	mdContext := metadata.AppendToOutgoingContext(
		context.Background(),
		"clientid", robotID, // 节点id
		"terminal_type", "webshell95", // 终端类型
		"session_id", sessionID, // 会话id
	)
	webshellCli, err := cli.MsgStreamWebMsgProxy(mdContext)
	if err != nil {
		conn.WriteMessage(1, []byte(err.Error()))
		return
	}
	// 新建会话
	ActiveWebShell(webshellCli, sessionID)

	// 结束连接
	connCtx, connCancel := context.WithCancel(context.Background())
	defer connCancel()

	// 连接断开处理
	conn.SetCloseHandler(func(code int, text string) error {
		connCancel()
		return nil
	})
	// 最近ping的时间
	latestReceiveWebMsgTime := time.Now().Unix()
	lastestPingTime := time.Now().Unix()
	latestNOPTime := time.Now().Unix()
	// 待发送的消息通道
	sendMsgChan := make(chan *devicec.Msg, 100)
	// 待接收的webMsg
	webMsgChan := make(chan []byte, 100)
	// 接收或发送消息
	var fileOffset int64
	go func() {
		defer func() {
			if err := recover(); err != nil {
				// 结束连接
				connCancel()
			}
		}()
		for {
			select {
			case <-connCtx.Done(): // 结束
				return
			case msg := <-sendMsgChan: // 向代理发送消息
				err := webshellCli.Send(msg)
				if err != nil {
					connCancel() // 结束连接
					logx.Error(context.Background(), "发送命令失败")
				}
			case data := <-webMsgChan: // 接收代理返回的消息，发送给web
				// 末尾检测
				for i := 0; i < 2; i++ {
					if !utf8.Valid(data) && len(data) > 1 {
						data = data[0 : len(data)-1]
					}
				}
				// 起始检测
				for i := 0; i < 2; i++ {
					if !utf8.Valid(data) && len(data) > 1 {
						data = data[i+1:]
					}
				}
				// utf8判断有效时，才发给前端
				if utf8.Valid(data) {
					err := conn.WriteMessage(1, data)
					if err != nil {
						connCancel() // 结束连接
						logx.Error(context.Background(), "回复web消息失败")
					}
				}
			}
		}
	}()
	// 定时发送NOP，会话保活
	go func() {
		tk := time.NewTicker(time.Duration(time.Second * 2))
		defer tk.Stop()
		for {
			select {
			// 会话结束
			case <-connCtx.Done():
				return
			// 每隔3秒定时发送
			case <-tk.C:
				// 保活
				webMsgChan <- []byte("")
				// NOP心跳
				nop := &devicec.WebShellNOP{}
				msg := &devicec.Msg{}
				msg.Name = string(nop.ProtoReflect().Descriptor().FullName())
				msg.Content, _ = proto.Marshal(nop)
				msg.SessionId = sessionID
				sendMsgChan <- msg
				// 判断ping超时
				if time.Now().Unix()-lastestPingTime > 60 {
					webMsgChan <- []byte("节点离线")
					go func() {
						time.Sleep(time.Second * 3)
						connCancel()
					}()
					logx.Info(context.Background(), "webshell ping timeout")
					return
				}
				// NOP超时
				if time.Now().Unix()-latestNOPTime > 30 {
					webMsgChan <- []byte("节点离线")
					go func() {
						time.Sleep(time.Second * 3)
						connCancel()
					}()
					logx.Info(context.Background(), "webshell nop timeout")
					return
				}
				// webscoket离线检测
				if time.Now().Unix()-latestReceiveWebMsgTime > 600 {
					webMsgChan <- []byte("操作超时")
					go func() {
						time.Sleep(time.Second * 3)
						connCancel()
					}()
					logx.Info(context.Background(), "webshell timeout")
					return
				}
			}
		}
	}()
	// 接收节点的消息，返回给web页面
	go func() {
		defer func() {
			if err := recover(); err != nil {
				// 结束连接
				connCancel()
			}
		}()
		for {
			select {
			// 结束
			case <-connCtx.Done():
				return
			default:
				msg, err := webshellCli.Recv()
				if err != nil {
					connCancel()
					return
				}
				switch msg.Name {
				// webshell data
				case string((&devicec.WebShellCMD{}).ProtoReflect().Descriptor().FullName()):
					cmd := &devicec.WebShellCMD{}
					if err := proto.Unmarshal(msg.Content, cmd); err == nil {
						webMsgChan <- cmd.Data
					}
				// ping回复
				case string((&devicec.Ping{}).ProtoReflect().Descriptor().FullName()):
					msg.Ack = 1
					sendMsgChan <- msg
					lastestPingTime = time.Now().Unix()
				// nop ack
				case string((&devicec.WebShellNOP{}).ProtoReflect().Descriptor().FullName()):
					latestNOPTime = time.Now().Unix()
				// 上传结果
				case string((&devicec.UploadFileRsp{}).ProtoReflect().Descriptor().FullName()):
					uploadFileRsp := &devicec.UploadFileRsp{}
					proto.Unmarshal(msg.Content, uploadFileRsp)
					res := make(map[string]interface{})
					res["code"] = uploadFileRsp.Code
					res["msg"] = uploadFileRsp.Msg
					if jsonstr, err := json.Marshal(res); err == nil {
						webMsgChan <- jsonstr
					}
				}
				time.Sleep(time.Millisecond * 50)
			}
		}
	}()
	for {
		select {
		case <-connCtx.Done():
			return
		default:
			// 读取信息，messageType: 1-文本，2-二进制
			_, message, err := conn.ReadMessage()
			if err == io.EOF {
				// 结束连接
				connCancel()
				return
			}
			if err == nil && len(message) > 0 {
				latestReceiveWebMsgTime = time.Now().Unix()
				var wsmsg struct {
					Type string `json:"type"`
					Cols int32  `json:"cols"`
					Rows int32  `json:"rows"`
				}
				// 判断消息类型
				if err := json.Unmarshal(message, &wsmsg); err == nil {
					// 调整虚拟终端的窗口大小
					if wsmsg.Type == "resize" {
						webshellSize := &devicec.WebShellSize{}
						webshellSize.Cols = wsmsg.Cols
						webshellSize.Rows = wsmsg.Rows
						content, _ := proto.Marshal(webshellSize)
						msg := &devicec.Msg{}
						msg.Name = string(webshellSize.ProtoReflect().Descriptor().FullName())
						msg.Content = content
						msg.SessionId = sessionID
						sendMsgChan <- msg
					}
				} else {
					// 上传文件的处理
					if dstDir != "" {
						msg := &devicec.Msg{}
						uploadFileReq := &devicec.UploadFileReq{}
						uploadFileReq.DstDir = dstDir
						uploadFileReq.FileName = fileName
						uploadFileReq.FileSize = fileSize
						uploadFileReq.FileMd5 = fileMD5
						uploadFileReq.Content = message
						uploadFileReq.Offset = fileOffset
						fileOffset += int64(len(message))
						content, _ := proto.Marshal(uploadFileReq)
						msg.Name = string(uploadFileReq.ProtoReflect().Descriptor().FullName())
						msg.Content = content
						msg.SessionId = sessionID
						sendMsgChan <- msg
						// 普通shell命令的处理
					} else {
						SaveShellRecord(robotID, userID, message)
						msg := &devicec.Msg{}
						cmd := &devicec.WebShellCMD{}
						cmd.Data = message
						content, _ := proto.Marshal(cmd)
						msg.Name = string(cmd.ProtoReflect().Descriptor().FullName())
						msg.Content = content
						msg.SessionId = sessionID
						sendMsgChan <- msg
					}
				}

			}
			time.Sleep(time.Millisecond * 50)
		}
	}
}

// ActiveWebShell 启动webshell
func ActiveWebShell(cli devicec.DeviceService_MsgStreamWebMsgProxyClient, sessionID string) {
	msg := &devicec.Msg{}
	cmd := &devicec.NewWebShellSession{}
	cmd.InitCmd = []string{
		"cd ~",                      // 进入默认工作目录
		"export LC_ALL=en_US.UTF-8", // 设置编码
	}
	content, _ := proto.Marshal(cmd)
	msg.Name = string(cmd.ProtoReflect().Descriptor().FullName())
	msg.Content = content
	msg.SessionId = sessionID
	cli.Send(msg)
}

var shellStr strings.Builder

// SaveShellRecord 保存shell操作记录
func SaveShellRecord(robotID string, userID string, data []byte) {
	// 有效的字符集合
	validStr := "1234567890abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ~!@#$%^&*()_+{}|:\"<>?`-=[]\\;',./\x20"
	for _, b := range data {
		if strings.Contains(validStr, string(b)) {
			shellStr.WriteByte(b)
		} else if b == []byte("\r")[0] {
			// 保存shell记录
			dao.WebshellLog{}.Add(context.Background(), dal.Use(mysql.NewDB()), models.WebshellLog{
				RobotID: robotID,
				UserID:  userID,
				Content: shellStr.String(),
			})
			// 重置shell内容
			shellStr.Reset()
		}
	}
}
